package cn.edu.flink.scala.tutorial.state.keyedstate

import cn.edu.flink.scala.tutorial.source.{ClickSource, Event}
import org.apache.flink.api.common.state.{MapState, MapStateDescriptor}
import org.apache.flink.streaming.api.functions.{KeyedProcessFunction, ProcessFunction}
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

import java.util.Map

object MapStateTest2 {
  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()

    val clickStream = env.addSource(new ClickSource).assignAscendingTimestamps(_.timestamp)

    // global
    clickStream.keyBy(x => true)
      .process(new KeyedProcessFunction[Boolean, Event, String] {

        val windowSize = 10000L
        lazy val countMapState: MapState[String, Long] = getRuntimeContext.getMapState(new MapStateDescriptor[String, Long]("countKV", classOf[String], classOf[Long]))

        override def processElement(value: Event, ctx: KeyedProcessFunction[Boolean, Event, String]#Context, out: Collector[String]): Unit = {
          countMapState.put(value.url, countMapState.get(value.url) + 1)

          val windowStart = value.timestamp / windowSize * windowSize
          val windowEnd = windowStart + windowSize

          ctx.timerService().registerEventTimeTimer(windowEnd - 1)
        }

        override def onTimer(timestamp: Long, ctx: KeyedProcessFunction[Boolean, Event, String]#OnTimerContext, out: Collector[String]): Unit = {
          import scala.collection.JavaConversions._

          for(count <- countMapState.iterator()) {
            out.collect(s"Window:${timestamp + 1 - windowSize} - ${timestamp + 1}, " +
              s"url: ${count.getKey}, count: ${count.getValue}")
          }
          countMapState.clear()
        }
      }).print()

    env.execute("MapStateTest2")
  }
}
